package leetcode_top;
import java.util.*;
import org.junit.*;

import utils.BracketUtils;
public class Ex378 {
    class Solution1 {
        public int kthSmallest(int[][] matrix, int k) {
            int m = matrix.length, n = matrix[0].length;
            if (m * n < k || k <= 0) return -1;
            Queue<Node> queue = new PriorityQueue<>((n1, n2) -> n1.val - n2.val);
            for (int i = 0; i < m; i++) {
                queue.offer(new Node(i, matrix[i][0], 0));
            }
            int cur = 0;
            Node node = null;
            while (cur < k) {
                node = queue.poll();
                node.pos++;
                cur++;
                if (node.pos == n) continue;
                if (cur == k) return node.val;
                node.val = matrix[node.idx][node.pos];
                queue.offer(node);
            }
            return -1;
        }
    
        class Node {
            int idx;
            int val;
            int pos;
            public Node(int idx, int val, int pos) {
                this.idx = idx;
                this.val = val;
                this.pos = pos;
            }

            public String toString() {
                return idx + "" + pos + "" + val;
            }
        }
    }

    class Solution {
        int[][] matrix;
        int m, n;
        public int kthSmallest(int[][] matrix, int k) {
            m = matrix.length;
            n = matrix[0].length;
            this.matrix = matrix;
            if (m * n < k || k <= 0) return -1;
            int lo = matrix[0][0], hi = matrix[m - 1][n - 1];
            int mid, num = 0;
            int[] pos = new int[m];
            while (lo < hi) {
                mid = (hi - lo) / 2 + lo;
                num = findMid(pos, mid);
                if (num >= k) {
                    hi = mid;
                } else {
                    lo = mid + 1;
                }
            }
            return lo;
        }

        public int findMid(int[] pos, int mid) {
            int count = 0;
            for (int i = 0; i < m; i++) {
                while (pos[i] >= 0 && matrix[i][pos[i]] > mid) {
                    pos[i]--;
                }
                while (pos[i] < n - 1 && matrix[i][pos[i] + 1] <= mid) {
                    pos[i]++;
                }
                count += (pos[i] + 1);
            }
            return count;
        }
    }

    @Test
    public void test() {
        Solution s = new Solution();
        int[][] nums = BracketUtils.to2DArray("[[1,2],[1,3]]");
        System.out.println(s.kthSmallest(nums, 2));
    }
}
